import pylsl
import mne
import numpy
import time
import threading
class LSL_MNE:
    def __init__(self,  \
                 max_Time=12000000, \
                 unit='uV', \
                 channels=[0, 1, 2, 3, 4, 5, 6, 7, 8], \
                 ch_names=['Fp1', 'Fp2', 'O2', 'Oz', 'Fz', 'Pz','Fpz','O1','STIM_014'], \
                 ch_types=['eeg', 'eeg', 'eeg', 'eeg', 'eeg', 'eeg','eeg','eeg', 'stim'], \
                 freq=250, \
                 lslID1='type', \
                 lslID2='eeg',
                 cache_save_Interval=30.0,\
                 save_Cache=True,
                 fileName='',\
                 have_Trigger=True
                 ):
        '''

        :param max_Time: 最大队列缓存时间（秒数）
        :param unit: 计算中采用的电压单位
        :param channels: 通道ID
        :param ch_names:    通道名
        :param ch_types:    通道类型
        :param freq:    采样频率
        :param lslID1:  输出属性1
        :param lslID2:  输出属性2
        :param cache_save_Interval: 缓存存储间隔
        :param save_Cache:  是否输出缓存
        '''
        self.ch_names=ch_names
        self.types=ch_types
        self.freq=freq
        self.max_Pendings=int(max_Time*60*self.freq)
        self.channels=channels
        self.stream=pylsl.resolve_stream(lslID1,lslID2)
        self.inlet=pylsl.stream_inlet(self.stream[0])
        self.interval=cache_save_Interval
        self.saveFlag=save_Cache
        self.fileName=fileName
        self.have_Trigger=have_Trigger
        #数据缓存
        self.dataTotal = []
        self.t_List=[]
        '''mne相关'''
        # Annotation
        self.onsets = []
        self.duration = []
        self.description = []
        self.unit=unit
        #   Events
        self.event_List = []

        #线程退出机制变量
        self.alive=True
    def Start_Streaming(self):
        t_Aq=threading.Thread(target=self.t_LSL,args=())
        t_Aq.setDaemon(True)
        t_Aq.start()
        if self.saveFlag:
            threading.Timer(interval=self.interval,function=self.t_DumpFile).start()
    def t_LSL(self):
        '''
            独立线程，循环获取LSL数据流
        '''
        if self.have_Trigger:
            dLen=len(self.channels)
        else:
            dLen = len(self.channels)-1
        while self.alive:
            dat,stamp=self.inlet.pull_sample()
            dat_Crop=dat[0:dLen]        #截取LSL数据以匹配预设电极数量
            #print(dat_Crop)
            # if self.unit=='uV':
            #     for i in range(dLen-1):
            #         dat_Crop[i]/=1000000.0#更换单位
            if self.have_Trigger:
                dat_Crop.append(0)          #添加STIM 通道
            self.dataTotal.append(dat_Crop)
            self.t_List.append(stamp)
            if (len(self.dataTotal)) > self.max_Pendings:
                '''采集数据流过长，丢弃前面的数据'''
                #   self.event_List.pop(0)
                self.t_List.pop(0)
                self.dataTotal.pop(0)
    def Set_Event(self,eid=1):
        '''
            在STIM通道配置Trigger   
            方法为把当前点的数据设置为eID的值
        '''
        if not self.have_Trigger:
            raise Exception("Unable to Set Trigger, Please set program's have_Trigger param to True")
        print('Setting Event:%d @ Data %d'%(eid,len(self.dataTotal)))
        self.dataTotal[-1][-1]=eid
    def t_DumpFile(self):
        '''
            开启定时缓存
        '''
        if self.alive:
            self.Dump2File()
            threading.Timer(interval=self.interval, function=self.t_DumpFile).start()
    def StopRecording(self):
        self.alive=False
    def Dump2File(self):
        '''构造MNE数据结构'''
        ''' 配置Info'''
        info = mne.create_info(ch_names=self.ch_names, sfreq=self.freq, ch_types=self.types)
        montage = mne.channels.make_standard_montage('standard_1020')
        info.set_montage(montage, on_missing='ignore')
        '''配置数据'''
        mne_Matrix=numpy.array(self.dataTotal).transpose()
        mne_raw = mne.io.RawArray(mne_Matrix, info)
        ''' 存储  '''
        mne_raw.save(self.fileName,overwrite=True)

if __name__=='__main__':
    l=LSL_MNE()
    l.Start_Streaming()
    #while len(l.dataTotal)==0: pass
    while True:
        time.sleep(10)